#include "style.h"

void EnergyLossCorrection(const Int_t ifile=0, const Int_t nevent = 100, const Bool_t kERatio=kTRUE, TCut cutEta="abs(Tracks[].fP[3])>0.5"){
  //
  // Draw difference between the BB energy loss estimated using the momenta at the vertex and 
  // more apropriate  momenta estimated at the inner wall of the TPC
  //
  // Effect due energy loss between the vertex and TPC can be described  using rough parameterization
  // p*= p+dEdx*[0] where the [0] is parameter describing the material between vertex and TPC
  // See the ratio between expected BetheBloch using moment a at vertex rsp. momenta at TPC inner wall
  // fbbb  for protons  and fbkk  for Kaons
  //
  //TF1 fbbMass("fbbMass","AliMathBase::BetheBlochAleph((x-AliMathBase::BetheBlochAleph(x/[1])*[0])/[1])/AliMathBase::BetheBlochAleph(x/[1])",0.5,2);
  //
  // Experimentally we got the correction coefficient which describe the material budget:
  // proton ~ 0.011
  // kaon   ~ 0.0075
  // pion   ~ 0.0062
  //
  // =================================================================================
  // =================================================================================
  //
  //Test by X.-G. Lu,  Di 11. Sep 22:22:15 CEST 2012
  //1) "/hera/alice/alien/alice/data/2011/LHC11h_2/000170593/ESDs/pass2/11000170593001.94/root_archive.zip#AliESDs.root"
  //proton:     1  p0           6.24232e-03   1.13336e-04   4.05017e-07   8.60577e-02
  //kaon:       1  p0           6.14199e-03   1.52037e-04   4.49197e-07  -2.08468e+00
  //pion:       1  p0           6.29124e-03   5.48414e-05   2.08836e-07  -3.87375e+00
  //2) "/hera/alice/alien/alice/data/2010/LHC10e/000130795/ESDs/pass2/10000130795042.1790/root_archive.zip#AliESDs.root"
  //proton:     1  p0           6.64228e-03   7.48949e-05   4.67794e-07   1.06520e+00
  //kaon:       1  p0           6.82135e-03   1.26729e-04   7.43744e-07  -2.42477e+00
  //pion:       1  p0           6.27682e-03   1.50058e-04   2.77864e-07  -4.45267e-01
  //3) "/hera/alice/alien/alice/data/2010/LHC10h/000138275/ESDs/pass2/10000138275027.430/root_archive.zip#AliESDs.root"
  //proton:     1  p0           6.09722e-03   9.92130e-05   4.04529e-07   9.52844e-02
  //kaon:       1  p0           6.21143e-03   1.20332e-04   4.11760e-07  -2.77414e+00
  //pion:       1  p0           7.77162e-03   4.15315e-05   1.59739e-06  -1.36363e+01
  //

  SetGlobalStyle(0);
  
  c1=new TCanvas;
 
  PadSetup(c1);
  c1->SetLeftMargin(0.14);
  c1->SetRightMargin(0.03);

  //----------

  TString fname;
  switch(ifile){
  case 0:
    fname = "/hera/alice/alien/alice/data/2011/LHC11h_2/000170593/ESDs/pass2/11000170593001.94/root_archive.zip#AliESDs.root";
    break;
  case 1:
    fname = "/hera/alice/alien/alice/data/2010/LHC10e/000130795/ESDs/pass2/10000130795042.1790/root_archive.zip#AliESDs.root";
    break;
  case 2:
    fname = "/hera/alice/alien/alice/data/2010/LHC10h/000138275/ESDs/pass2/10000138275027.430/root_archive.zip#AliESDs.root";
    break;
  default:
    cout<<"wrong ifile"<<endl;
    return;
  }

  TFile *fesd = TFile::Open(fname);
  TTree * esdTree = fesd->Get("esdTree");
  esdTree->SetCacheSize(1000000000);
  const TCut cutNcl = "Tracks[].fTPCncls>70&&Tracks[].fITSncls>3";  
  esdTree->Draw("AliExternalTrackParam::BetheBlochSolid(Tracks[].P()/0.9)/AliExternalTrackParam::BetheBlochSolid(Tracks[].fIp.P()/0.9):Tracks[].P()",cutNcl+cutEta+"Tracks[].P()<2","",nevent);

  //should have the direct correspondance between x ~ Tracks[].P and sqrt(TMath::Power(sqrt(x*x+[1]*[1])-[0]*AliExternalTrackParam::BetheBlochSolid(x/[1]),2)-TMath::Power([1],2)) ~ Tracks[].fIp.P()
  TF1 *fbbMass=0x0;
  if(kERatio){
    fbbMass = new TF1("fbbMass","AliExternalTrackParam::BetheBlochSolid(x/[1])/AliExternalTrackParam::BetheBlochSolid(TMath::Sqrt(TMath::Power(TMath::Sqrt(x*x+[1]*[1])-[0]*AliExternalTrackParam::BetheBlochSolid(x/[1]),2)-TMath::Power([1],2))/[1])",0.5,5);
  }
  else{ 
    fbbMass = new TF1("fbbMass","x/TMath::Sqrt(TMath::Power(TMath::Sqrt(x*x+[1]*[1])-[0]*AliExternalTrackParam::BetheBlochSolid(x/[1]),2)-TMath::Power([1],2))",0.5,5);
  }

  //
  // ratio BB if we use the momenta from vertex in resp. to the BB with proper momenta in the TPC
  //
  const Double_t mass[]={0.938272, 0.139570, 0.493677, 5.109989e-04 };
  const char * pname[]={"Proton","Pion","Kaon","Electron"};
  const Int_t aliid[]={AliPID::kProton,AliPID::kPion,AliPID::kKaon,AliPID::kElectron};
  const Double_t pmin[]={0.5,0.3,0.4,0.3};

  for(Int_t ii=0; ii<3; ii++){
    TString var;
     if(kERatio){
       var = Form("AliExternalTrackParam::BetheBlochSolid(Tracks[].P()/%e)/AliExternalTrackParam::BetheBlochSolid(Tracks[].fIp.P()/%e):Tracks[].P()>>prof%s(40, %f,3)",mass[ii],mass[ii], pname[ii], pmin[ii]);
     }
     else{
       var = Form("Tracks[].P()/Tracks[].fIp.P():Tracks[].P()>>prof%s(40, %f,3)",pname[ii], pmin[ii]);
     }
     const TString icut(Form("Tracks[].fTOFr[%d]*Tracks[].fITSr[%d]>0.1",aliid[ii], aliid[ii]));
     printf("check %d: %s %s\n", ii, var.Data(), icut.Data());
     TCut tmpcut(icut);
     esdTree->Draw(var, cutNcl+cutEta+tmpcut,"prof",nevent);  
  }

  profPion->SetMarkerColor(1);
  profPion->SetLineColor(1);
  profPion->SetMarkerStyle(20);

  profProton->SetMarkerColor(4);
  profProton->SetLineColor(4);
  profProton->SetMarkerStyle(30);

  profKaon->SetMarkerColor(2);
  profKaon->SetLineColor(2);
  profKaon->SetMarkerStyle(22);

  if(kERatio){
    profPion->SetMinimum(0.8);
    profPion->SetMaximum(1.02);
  }
  else{
    profPion->SetMinimum(0.98);
    profPion->SetMaximum(1.2);
  }
  profPion->GetXaxis()->SetTitle("p_{vertex} (GeV/c)"); 
  profPion->GetYaxis()->SetTitle(kERatio?"dEdx(p_{vertex})/dEdx(p_{TPC})":"p_{vertex}/p_{inner TPC}"); 
  gPad->SetLogx();
  
  const TDatabasePDG *ppdg= TDatabasePDG::Instance();
  fbbMass->SetParameters(0,ppdg->GetParticle("proton")->Mass());
  fbbMass->FixParameter(1,ppdg->GetParticle("proton")->Mass());
  fbbMass->SetLineColor(profProton->GetMarkerColor());
  fp=(TF1*)fbbMass->Clone("fp");
  printf("fitting proton...\n");
  profProton->Fit(fp,"","", 0.4,3);
  
  fbbMass->SetParameters(0,ppdg->GetParticle("K+")->Mass());
  fbbMass->FixParameter(1,ppdg->GetParticle("K+")->Mass());
  fbbMass->SetLineColor(profKaon->GetMarkerColor());
  fK=(TF1*)fbbMass->Clone("fK");
  printf("fitting kaon...\n");
  profKaon->Fit(fK,"","", 0.3,3);
  
  fbbMass->SetParameters(0,ppdg->GetParticle("pi+")->Mass());
  fbbMass->FixParameter(1,ppdg->GetParticle("pi+")->Mass());
  fbbMass->SetLineColor(profPion->GetMarkerColor());
  fpi=(TF1*)fbbMass->Clone("fpi");
  printf("fitting pion...\n");
  profPion->Fit(fpi,"","",0.2,3);
  
  ResetStyle(profPion);
  ResetStyle(profKaon);
  ResetStyle(profProton);
  profPion->Draw();
  profKaon->Draw("same");
  profProton->Draw("same");
  fK->Draw("same");
  fpi->Draw("same");
  fp->Draw("same");
  
  const TString tag=fname(34,6);
  TLegend * legend = new TLegend(0.7,kERatio ? 0.2 : 0.7 ,0.85, kERatio ? 0.35: 0.85);
  ResetStyle(legend,0.15);
  legend->SetHeader(tag);
  legend->AddEntry(profPion,"Pion");
  legend->AddEntry(profKaon,"Kaon");
  legend->AddEntry(profProton,"Proton");
  legend->Draw();

  c1->Print(tag+"_energyLossCorr.eps");
}

//oct 10, 2012, 15:59
/*
root [1] EnergyLossCorrection(1, 1000000, 1)
Warning in <TClass::TClass>: no dictionary for class AliESDCentrality is available
check 0: AliExternalTrackParam::BetheBlochSolid(Tracks[].P()/9.382720e-01)/AliExternalTrackParam::BetheBlochSolid(Tracks[].fIp.P()/9.382720e-01):Tracks[].P()>>profProton(40, 0.500000,3) Tracks[].fTOFr[4]*Tracks[].fITSr[4]>0.1
check 1: AliExternalTrackParam::BetheBlochSolid(Tracks[].P()/1.395700e-01)/AliExternalTrackParam::BetheBlochSolid(Tracks[].fIp.P()/1.395700e-01):Tracks[].P()>>profPion(40, 0.300000,3) Tracks[].fTOFr[2]*Tracks[].fITSr[2]>0.1
check 2: AliExternalTrackParam::BetheBlochSolid(Tracks[].P()/4.936770e-01)/AliExternalTrackParam::BetheBlochSolid(Tracks[].fIp.P()/4.936770e-01):Tracks[].P()>>profKaon(40, 0.400000,3) Tracks[].fTOFr[3]*Tracks[].fITSr[3]>0.1
fitting proton...
 FCN=158.065 FROM MIGRAD    STATUS=CONVERGED      19 CALLS          20 TOTAL
                     EDM=1.72263e-10    STRATEGY= 1      ERROR MATRIX ACCURATE 
  EXT PARAMETER                                   STEP         FIRST   
  NO.   NAME      VALUE            ERROR          SIZE      DERIVATIVE 
   1  p0           4.07172e+00   7.10638e-02   4.37627e-04   2.61193e-04
   2  p1           9.38272e-01     fixed    
fitting kaon...
 FCN=37.4516 FROM MIGRAD    STATUS=CONVERGED      19 CALLS          20 TOTAL
                     EDM=1.30024e-10    STRATEGY= 1      ERROR MATRIX ACCURATE 
  EXT PARAMETER                                   STEP         FIRST   
  NO.   NAME      VALUE            ERROR          SIZE      DERIVATIVE 
   1  p0           4.01534e+00   8.13919e-02   2.46438e-04   1.98128e-04
   2  p1           4.93677e-01     fixed    
fitting pion...
 FCN=94.9235 FROM MIGRAD    STATUS=CONVERGED      18 CALLS          19 TOTAL
                     EDM=7.53239e-13    STRATEGY= 1      ERROR MATRIX ACCURATE 
  EXT PARAMETER                                   STEP         FIRST   
  NO.   NAME      VALUE            ERROR          SIZE      DERIVATIVE 
   1  p0           4.20545e+00   4.07242e-02   1.94752e-04  -3.01390e-05
   2  p1           1.39570e-01     fixed    
Info in <TCanvas::Print>: eps file LHC10e_energyLossCorr.eps has been created
*/
